
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class ResidualBlock(nn.Module):
    def __init__(self, channels=64):
        super().__init__()
        self.conv1 = nn.Conv2d(channels, channels, 3, 1, 1)
        self.conv2 = nn.Conv2d(channels, channels, 3, 1, 1)

    def forward(self, x):
        return x + self.conv2(F.relu(self.conv1(x)))

class Encoder(nn.Module):
    def __init__(self, in_ch=3, channels=64, num_blocks=8):
        super().__init__()
        layers = [nn.Conv2d(in_ch, channels, 3, 1, 1)]
        for _ in range(num_blocks):
            layers.append(ResidualBlock(channels))
        self.body = nn.Sequential(*layers)

    def forward(self, x):
        return self.body(x)

def build_cov_from_params(sigx, sigy, rho, eps=1e-6):
    sigx = F.softplus(sigx) + eps
    sigy = F.softplus(sigy) + eps
    rho = torch.tanh(rho).clamp(-0.999, 0.999)
    a = sigx * sigx
    b = rho * sigx * sigy
    c = sigy * sigy
    det = a * c - b * b + eps
    inv = torch.stack([torch.stack([ c/det, -b/det], dim=-1),
                       torch.stack([-b/det,  a/det], dim=-1)], dim=-2)
    return inv, det

def gaussian_weight(query_xy, mean_xy, inv_cov):
    
    breakpoint()
    d = (query_xy - mean_xy)[..., None, :]
    md2 = torch.matmul(torch.matmul(d, inv_cov), d.transpose(-1,-2)).squeeze(-1).squeeze(-1)
    return torch.exp(-0.5 * md2)

def splat_2d_gaussians(hr_H, hr_W, coords, colors, sigx, sigy, rho, opacity=None, hr_to_lr_scale=1.0, mode="sum", device=None):
    if coords.dim() == 2:
        coords = coords.unsqueeze(0)
        colors = colors.unsqueeze(0)
        sigx = sigx.unsqueeze(0); sigy = sigy.unsqueeze(0); rho = rho.unsqueeze(0)
        if opacity is not None:
            opacity = opacity.unsqueeze(0)
    B, G, _ = coords.shape
    C = colors.shape[-1]
    device = device or colors.device

    ys = torch.arange(hr_H, device=device) + 0.5
    xs = torch.arange(hr_W, device=device) + 0.5
    yy, xx = torch.meshgrid(ys, xs, indexing="ij")
    grid = torch.stack([xx / hr_to_lr_scale, yy / hr_to_lr_scale], dim=-1).view(1, 1, hr_H*hr_W, 2)
    grid = grid.expand(B, G, -1, -1)

    breakpoint()
    
    inv, det = build_cov_from_params(sigx, sigy, rho)
    if opacity is None:
        alpha_i = torch.ones(B, G, device=device)
    else:
        alpha_i = torch.sigmoid(opacity)

    means = coords.unsqueeze(2)
    inv_cov = inv.unsqueeze(2)
    w = gaussian_weight(grid, means, inv_cov)
    
    if mode == "sum":
        coeff = (alpha_i[..., None]) * w
        colored = colors[..., None] * coeff.unsqueeze(2)
        out = colored.sum(dim=1).view(B, C, hr_H, hr_W)
        return out
    elif mode == "alpha":
        a = (alpha_i[..., None]) * w
        T = torch.cumprod(1.0 - a + 1e-6, dim=1)
        one = torch.ones_like(T[:, :1, :])
        T_prev = torch.cat([one, T[:, :-1, :]], dim=1)
        weights = a * T_prev
        colored = colors[..., None] * weights.unsqueeze(2)
        out = colored.sum(dim=1).view(B, C, hr_H, hr_W)
        return out
    else:
        raise ValueError("mode must be 'sum' or 'alpha'")

class GaussianBank(nn.Module):
    def __init__(self, num_gaussians=100, init_sigma=0.8, init_rho=0.0, init_opacity=0.0):
        super().__init__()
        self.num_gaussians = num_gaussians
        self.sigx = nn.Parameter(torch.full((num_gaussians,), math.log(math.exp(init_sigma)-1.0)))
        self.sigy = nn.Parameter(torch.full((num_gaussians,), math.log(math.exp(init_sigma)-1.0)))
        self.rho  = nn.Parameter(torch.full((num_gaussians,), init_rho))
        self.opa  = nn.Parameter(torch.full((num_gaussians,), init_opacity))

    def forward(self):
        return {
            "sigx": self.sigx,
            "sigy": self.sigy,
            "rho":  self.rho,
            "opacity": self.opa,
        }

class SGSRenderer(nn.Module):
    def __init__(self, gaussian_bank: GaussianBank, temperature=1.0, use_alpha_sum=True):
        super().__init__()
        self.bank = gaussian_bank
        self.temperature = temperature
        self.use_alpha_sum = use_alpha_sum

    @staticmethod
    def pixel_centers(H, W, device):
        ys = torch.arange(H, device=device) + 0.5
        xs = torch.arange(W, device=device) + 0.5
        yy, xx = torch.meshgrid(ys, xs, indexing="ij")
        return torch.stack([xx, yy], dim=-1)

    def forward(self, feat_lr, logits, hr_size, scale, train_mode=True):
        B, C, H, W = feat_lr.shape
        G = logits.shape[1]
        Hh, Wh = hr_size
        device = feat_lr.device

        sigx, sigy, rho, opa, color = torch.split(logits, [1,1,1,1,3], dim=1)
        
        
        
        centers = self.pixel_centers(H, W, device).view(1, H*W, 2).expand(B,-1,-1)
        breakpoint()
        colors = feat_lr.flatten(2,3).transpose(1,2)

        rendered = splat_2d_gaussians(
            Hh, Wh,
            coords=centers, colors=colors,
            sigx=sigx, sigy=sigy, rho=rho, opacity=opa,
            hr_to_lr_scale=scale,
            mode="sum" if self.use_alpha_sum else "alpha",
            device=device
        )
        return rendered


class Decoder(nn.Module):
    def __init__(self, in_ch=64, out_ch=3):
        super().__init__()
        self.tail = nn.Sequential(
            nn.Conv2d(in_ch, in_ch, 3, 1, 1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_ch, out_ch, 3, 1, 1),
            nn.Sigmoid(),
        )

    def forward(self, x):
        return self.tail(x)


class GaussianSRModel(nn.Module):
    def __init__(self, in_ch=3, channels=64, num_blocks=2, num_gaussians=100, temperature=1.0, use_alpha_sum=True):
        super().__init__()
        self.encoder = Encoder(in_ch=in_ch, channels=channels, num_blocks=num_blocks)
        self.classifier = nn.Conv2d(channels, 7, 1)
        self.bank = GaussianBank(num_gaussians=num_gaussians)
        self.renderer = SGSRenderer(self.bank, temperature=temperature, use_alpha_sum=use_alpha_sum)
        self.decoder = Decoder(in_ch=channels, out_ch=3)

    def forward(self, lr, hr_size, scale, train_mode=True):
        feat = self.encoder(lr)
        logits = self.classifier(feat)
        feat_hr = self.renderer(feat, logits, hr_size=hr_size, scale=scale, train_mode=train_mode)
        sr = self.decoder(feat_hr)
        return sr, feat_hr, logits


